package io.gitee.devaper.starlight.websocket

import io.gitee.devaper.starlight.common.constant.*
import org.springframework.aop.support.AopUtils
import org.springframework.context.*
import org.springframework.context.annotation.*
import org.springframework.core.type.AnnotationMetadata
import org.springframework.web.socket.WebSocketHandler
import org.springframework.web.socket.config.annotation.*
import org.springframework.web.socket.server.HandshakeHandler
import org.springframework.web.socket.server.HandshakeInterceptor

/**
 * WebSocketController的注册器, 自动扫描@WebSocketController标记的类并注册
 */
@Configuration(proxyBeanMethods = false)
@EnableWebSocket
open class WebSocketControllerRegistrar : ApplicationContextAware, WebSocketConfigurer {

    private lateinit var applicationContext: ApplicationContext

    @Suppress("unchecked_cast")
    override fun registerWebSocketHandlers(registry: WebSocketHandlerRegistry) {

        applicationContext.getBeansWithAnnotation(WebSocketController::class.java).values.forEach { bean ->
            if (bean !is WebSocketHandler) {
                throw ClassCastException("${AopUtils.getTargetClass(bean).canonicalName} -> ${WebSocketHandler::class.java.canonicalName}")
            }
            // 解析注解
            val annotationMetadata      = AnnotationMetadata.introspect(AopUtils.getTargetClass(bean))
            val annotationAttributes    = annotationMetadata.getAnnotationAttributes(WebSocketController::class.java.canonicalName)!!

            // 取出注解属性
            val paths                   = annotationAttributes[WebSocketController::paths.name] as Array<String>
            val allowedOrigins          = annotationAttributes[WebSocketController::allowedOrigins.name] as Array<String>
            val allowedOriginPatterns   = annotationAttributes[WebSocketController::allowedOriginPatterns.name] as Array<String>
            val interceptors            = annotationAttributes[WebSocketController::interceptors.name] as Array<Class<HandshakeInterceptor>>
            val handshakeHandler        = annotationAttributes[WebSocketController::handshakeHandler.name] as Array<Class<HandshakeHandler>>

            // 注册WebSocketHandler
            val registration = registry.addHandler(bean, *(paths.map { WS_URL_PREFIX + it }.toTypedArray()))
            if (allowedOrigins.isNotEmpty()) {
                registration.setAllowedOrigins(*allowedOrigins)
            }
            if (allowedOriginPatterns.isNotEmpty()) {
                registration.setAllowedOriginPatterns(*allowedOriginPatterns)
            }
            if (interceptors.isNotEmpty()) {
                val interceptorBeans = ArrayList<HandshakeInterceptor>(interceptors.size)
                interceptors.forEach { interceptor ->
                    interceptorBeans.add(applicationContext.getBean(interceptor))
                }
                registration.addInterceptors(*interceptorBeans.toTypedArray())
            }
            // 如果Bean本身也实现了HandshakeInterceptor接口, 那么将它自己也作为一个Interceptor
            if (bean is HandshakeInterceptor) {
                registration.addInterceptors(bean)
            }
            if (handshakeHandler.isNotEmpty()) {
                registration.setHandshakeHandler(applicationContext.getBean(handshakeHandler.first()))
            } else if (bean is HandshakeHandler) {
                // 如果Bean本身也实现了HandshakeHandler接口, 那么将它自己作为HandshakeHandler
                registration.setHandshakeHandler(bean)
            }
        }
    }

    override fun setApplicationContext(applicationContext: ApplicationContext) {
        this.applicationContext = applicationContext;
    }

}